#include <stdio.h>
#include <stdlib.h>
#include <libavformat/avformat.h>

#include "util.h"
#include "audiosource.h"
#include "audioencoder_av.h"

struct AUDIOENCODER_AV {
	AUDIOENCODER encoder;

	AUDIOSOURCE *source;
	long framepos;

    AVFormatContext *oc;
    AVStream *audio_st;

	int16_t *samples;
	int frame_size;

	float *float_samples;

/*	int init_frames_written;
	lame_global_flags *lame;*/
};

/**
 * Is format available.
 */
static int is_sample_fmt_available(const enum AVSampleFormat *fmt, int cand) {
	int i=0;

	while (fmt[i]!=-1) {
		if (fmt[i]==cand)
			return 1;

		i++;
	}

	return 0;
}

/**
 * Choose format.
 */
static int choose_sample_format(const enum AVSampleFormat *available) {
	if (is_sample_fmt_available(available,AV_SAMPLE_FMT_S16))
		return AV_SAMPLE_FMT_S16;

	if (is_sample_fmt_available(available,AV_SAMPLE_FMT_S16P))
		return AV_SAMPLE_FMT_S16P;

	if (is_sample_fmt_available(available,AV_SAMPLE_FMT_S32))
		return AV_SAMPLE_FMT_S32;

	if (is_sample_fmt_available(available,AV_SAMPLE_FMT_S32P))
		return AV_SAMPLE_FMT_S32P;

	if (is_sample_fmt_available(available,AV_SAMPLE_FMT_DBL))
		return AV_SAMPLE_FMT_DBL;

	if (is_sample_fmt_available(available,AV_SAMPLE_FMT_DBLP))
		return AV_SAMPLE_FMT_DBLP;

	if (is_sample_fmt_available(available,AV_SAMPLE_FMT_FLT))
		return AV_SAMPLE_FMT_FLT;

	if (is_sample_fmt_available(available,AV_SAMPLE_FMT_FLTP))
		return AV_SAMPLE_FMT_FLTP;

    int i=0;

    while (available[i]!=-1) {
    	TRACE("sampleformat %d: %d %s\n",i,available[i],av_get_sample_fmt_name(available[i]));
    	i++;
    }

	fail("no suitable sample format found...");
}

/*
 * add an audio output stream
 */
static AVStream *add_audio_stream(AVFormatContext *oc, enum AVCodecID codec_id)
{
    AVCodecContext *c;
    AVStream *st;
    AVCodec *codec;

    /* find the audio encoder */
    TRACE("finding encoder for: %d\n",codec_id);
    codec = avcodec_find_encoder(codec_id);

//	codec = avcodec_find_encoder_by_name("libvorbis");
//	codec = avcodec_find_encoder_by_name("vorbis");
    if (!codec) {
        fprintf(stderr, "codec not found\n");
        exit(1);
    }

    TRACE("got encoder: %s\n",codec->name);

    st = avformat_new_stream(oc, codec);
    if (!st) {
        fprintf(stderr, "Could not alloc stream\n");
        exit(1);
    }
    st->id = 1;

    c = st->codec;

    c->codec=codec;
    TRACE("default: bitrate: %d, samplerate: %d, fmt: %d\n",c->bit_rate,c->sample_rate,c->sample_fmt);

//    fail("exit");
//    codec->sample_fmts;

    /* put sample parameters */
    c->sample_fmt  = choose_sample_format(codec->sample_fmts);
//    c->sample_fmt  = AV_SAMPLE_FMT_S32P;
//    c->sample_fmt  = AV_SAMPLE_FMT_DBL;
//    c->bit_rate    = 64000;
    c->bit_rate    = 128000;
    c->sample_rate = 44100;
    c->channels    = 2;

    // some formats want stream headers to be separate
    if (oc->oformat->flags & AVFMT_GLOBALHEADER)
        c->flags |= CODEC_FLAG_GLOBAL_HEADER;

    return st;
}

/**
 * Open codec.
 */
static void open_audio(AUDIOENCODER_AV *encoder, AVFormatContext *oc, AVStream *st)
{
    AVCodecContext *c;

    c = st->codec;

    TRACE("opening audio codec: %p\n",c->codec);

    /* open it */
    if (avcodec_open2(c, NULL, NULL) < 0)
    	fail("could not open codec\n");

    if (c->codec->capabilities & CODEC_CAP_VARIABLE_FRAME_SIZE)
    	fail("didn't expect frame size to be variable.");

    if (c->channels!=2)
    	fail("expected 2 channels for output.");

//	TRACE("variable frame size: %d\n",c->codec->capabilities & CODEC_CAP_VARIABLE_FRAME_SIZE);

    TRACE("frame size: %d\n",c->frame_size);
	TRACE("sample fmt: %d\n",c->sample_fmt);
	//fail("exit");

/*	if (c->sample_fmt!=AV_SAMPLE_FMT_S16)
		fail("expected sample format to be AV_SAMPLE_FMT_S16.");*/

	encoder->frame_size=c->frame_size;
	encoder->samples = av_malloc(encoder->frame_size *
		av_get_bytes_per_sample(c->sample_fmt) *
		c->channels);

	encoder->float_samples=malloc(encoder->frame_size*c->channels*sizeof(float));

	if (!encoder->float_samples || !encoder->samples)
		fail("unable to malloc sample buffer");

/*
    t     = 0;
    tincr = 2 * M_PI * 110.0 / c->sample_rate;
    increment frequency by 110 Hz per second
    tincr2 = 2 * M_PI * 110.0 / c->sample_rate / c->sample_rate;

    if (c->codec->capabilities & CODEC_CAP_VARIABLE_FRAME_SIZE)
        audio_input_frame_size = 10000;
    else
        audio_input_frame_size = c->frame_size;
    samples = av_malloc(audio_input_frame_size *
                        av_get_bytes_per_sample(c->sample_fmt) *
                        c->channels);
    */
}

/**
 * Close audio.
 */
static void close_audio(AVFormatContext *oc, AVStream *st)
{
    avcodec_close(st->codec);

//    av_free(samples);
}

/**
 * End of file?
 */
static int eof(AUDIOENCODER *e) {
	AUDIOENCODER_AV *encoder=(AUDIOENCODER_AV *)e;

	if (encoder->framepos>=audiosource_get_numsamples(encoder->source))
		return 1;

	else
		return 0;
}

/**
 * Clamp float value.
 */
static float clamp(float v) {
	if (v>1.0)
		v=1.0;

	if (v<-1.0)
		v=-1.0;

	return v;
}

/**
 * Pack samples for encoding.
 */
static void pack_samples(void *target, enum AVSampleFormat format, float *samples, int num_samples) {
	int16_t *int16target=(int16_t *)target;
	float *floattarget=(float *)target;
	double *doubletarget=(double *)target;
	int32_t *int32target=(int32_t *)target;

	int i;
	for (i=0; i<num_samples; i++) {
		float left=clamp(samples[i*2]);
		float right=clamp(samples[i*2+1]);

		switch (format) {
			case AV_SAMPLE_FMT_S16:
				int16target[i*2]=left*0x7fff;
				int16target[i*2+1]=right*0x7fff;
				break;

			case AV_SAMPLE_FMT_S16P:
				int16target[i]=left*0x7fff;
				int16target[i+num_samples]=right*0x7fff;
				break;

			case AV_SAMPLE_FMT_S32:
				int32target[i*2]=left*0x7fffffff;
				int32target[i*2+1]=right*0x7fffffff;
				break;

			case AV_SAMPLE_FMT_S32P:
				int32target[i]=left*0x7fffffff;
				int32target[i+num_samples]=right*0x7fffffff;
				break;

			case AV_SAMPLE_FMT_FLT:
				floattarget[i*2]=left;
				floattarget[i*2+1]=right;
				break;

			case AV_SAMPLE_FMT_FLTP:
				floattarget[i]=left;
				floattarget[i+num_samples]=right;
				break;

			case AV_SAMPLE_FMT_DBL:
				doubletarget[i*2]=left;
				doubletarget[i*2+1]=right;
				break;

			case AV_SAMPLE_FMT_DBLP:
				doubletarget[i]=left;
				doubletarget[i+num_samples]=right;
				break;

			default:
				fail("sample format not supported.");
				break;
		}
	}
}

/**
 * Encode.
 */ 
static int encode(AUDIOENCODER *e, int length) {
	AUDIOENCODER_AV *encoder=(AUDIOENCODER_AV *)e;
    AVFrame *frame;
    AVCodecContext *c=encoder->audio_st->codec;

	int numframes=length*44100/1000;
    int startpos=encoder->framepos;
	int last=encoder->framepos+numframes;
    int got_packet;

	while (encoder->framepos<last) {
	    AVPacket pkt = { 0 }; // data and size must be 0;
		frame=avcodec_alloc_frame();
		av_init_packet(&pkt);

		// get samples and put in samples..	
		int res=audiosource_get_samples(encoder->source,encoder->float_samples,encoder->frame_size);
		if (!res)
			memset_float(encoder->float_samples,0,encoder->frame_size*2);

/*		for (i=0; i<encoder->frame_size*2; i++) {
			float sample=encoder->float_samples[i];

			if (sample>1.0)
				sample=1.0;

			if (sample<-1.0)
				sample=-1.0;

			encoder->samples[i]=sample*0x7fff;
		}*/

		pack_samples(encoder->samples,c->sample_fmt,encoder->float_samples,encoder->frame_size);

		/*for (i=0; i<encoder->frame_size; i++) {
			float sample;

			sample=encoder->float_samples[i*2];

			if (sample>1.0)
				sample=1.0;

			if (sample<-1.0)
				sample=-1.0;

			encoder->samples[i]=sample*0x7fff;

			sample=encoder->float_samples[i*2+1];

			if (sample>1.0)
				sample=1.0;

			if (sample<-1.0)
				sample=-1.0;

			encoder->samples[encoder->frame_size+i]=sample*0x7fff;
		}*/

		frame->nb_samples = encoder->frame_size;

		/*TRACE("buffer size: %d\n",av_samples_get_buffer_size(NULL,2,encoder->frame_size,AV_SAMPLE_FMT_S32,1));
		TRACE("buffer size: %d\n",av_samples_get_buffer_size(NULL,2,encoder->frame_size,AV_SAMPLE_FMT_S32P,1));
		TRACE("size: %d\n",encoder->frame_size *
			av_get_bytes_per_sample(c->sample_fmt) *
			c->channels);*/

		int fillres=avcodec_fill_audio_frame(frame, c->channels, c->sample_fmt,
			(uint8_t *)encoder->samples,
			encoder->frame_size *
			av_get_bytes_per_sample(c->sample_fmt) *
			c->channels, 1);

		if (fillres<0)
			fail("unable to fill audio frame: %d",fillres);

		int encres=avcodec_encode_audio2(c, &pkt, frame, &got_packet);
//		if (!got_packet)
		if (encres!=0)
			fail("got no packet, encres=%d",encres);

//		if (got_packet) {
			pkt.stream_index = encoder->audio_st->index;

			/* Write the compressed frame to the media file. */
			//TRACE("*********** writing: %d\n",length);
			if (av_interleaved_write_frame(encoder->oc, &pkt) != 0) {
				fail("unable to write audio frame.");
			}
//		}

		encoder->framepos+=encoder->frame_size;
	}

	return encoder->framepos-startpos;
}

/**
 * Destructor.
 */
static void dispose(AUDIOENCODER *e) {
	AUDIOENCODER_AV *encoder=(AUDIOENCODER_AV *)e;

	int i;

    av_write_trailer(encoder->oc);

    /* Close each codec. */
	close_audio(encoder->oc, encoder->audio_st);

    /* Free the streams. */
    for (i = 0; i < encoder->oc->nb_streams; i++) {
        av_freep(&encoder->oc->streams[i]->codec);
        av_freep(&encoder->oc->streams[i]);
    }

	avio_close(encoder->oc->pb);

    /* free the stream */
    av_free(encoder->oc);

	free(encoder);
}

/**
 * Constrcutor.
 */
AUDIOENCODER_AV *audioencoder_av_create(char *target, AUDIOSOURCE *source, int position, char *format) {
	AUDIOENCODER_AV *encoder;
	char *filename;

	encoder=nicemalloc(sizeof(AUDIOENCODER_AV));

	encoder->encoder.encode_func=encode;
	encoder->encoder.eof_func=eof;
	encoder->encoder.dispose_func=dispose;

	encoder->source=source;
	encoder->framepos=(long long)position*44100/1000;

	audiosource_set_streampos(encoder->source,encoder->framepos);

	filename=target;
	if (!filename)
		filename="pipe:1";

    av_register_all();

	avformat_alloc_output_context2(&encoder->oc, NULL, format, filename);
	if (!encoder->oc)
		fail("Unable to create output context.");

    AVOutputFormat *fmt=encoder->oc->oformat;

    TRACE("mp3 codec id: %d",AV_CODEC_ID_MP3);

/*    if (!strcmp(format,"ogg"))
		encoder->audio_st = add_audio_stream(encoder->oc, AV_CODEC_ID_LIBVORBIS);

	else*/
		encoder->audio_st = add_audio_stream(encoder->oc, fmt->audio_codec);

	open_audio(encoder, encoder->oc, encoder->audio_st);

	if (avio_open(&encoder->oc->pb, filename, AVIO_FLAG_WRITE) < 0)
		fail("Could not open output file.");

    avformat_write_header(encoder->oc, NULL);

	return encoder;
}
